package com.cloudea.learn.layer;

import com.cloudea.learn.Tensor;

public class ElementSub extends Layer {
    @Override
    protected Tensor forward(Context context, Tensor... inputs) {
        Tensor a = inputs[0];
        Tensor b = inputs[1];
        Tensor c = new Tensor(a.getShape());
        for (int i = 0; i < c.getSize(); i++){
            c.setByLocation(i, a.getByLocation(i) - b.getByLocation(i));
        }
        return c;
    }

    @Override
    public Tensor[] backward(Context context, Tensor grid, Tensor... inputs) {
        Tensor a = inputs[0];
        Tensor b = inputs[1];
        int n = a.getSize();
        Tensor gridA = new Tensor(n);
        Tensor gridB = new Tensor(n);
        for (int i = 0; i < n; i++){
            gridA.setByLocation(i, grid.getByLocation(i));
            gridB.setByLocation(i, - grid.getByLocation(i));
        }
        return new Tensor[]{gridA, gridB};
    }
}
